import torch.nn as nn
from torchcrf import CRF


class BERT_CRF(nn.Module):
    def __init__(self, pretrained_model, tag2id, crf=True):
        super(BERT_CRF, self).__init__()
        self.tag2id = tag2id
        self.pretrained_model = pretrained_model.from_pretrained('deep_learning/bert/save_model/bert_model_0', num_labels=len(self.tag2id))
        if crf:
            self.crf = CRF(len(self.tag2id))
    def forward(self, x):
        result = self.pretrained_model(x, token_type_ids=None, attention_mask=(x>0))
        outputs = result.logits
        #CRF
        outputs = self.crf.decode(outputs)
        return outputs
    def log_likelihood(self, x, tags):
        result = self.pretrained_model(x, token_type_ids=None, attention_mask=(x>0))
        outputs = result.logits
        #CRF
        return - self.crf(outputs, tags)